import cupy as cp
from enum import Enum


# 1. 定义操作枚举
class Operation(Enum):
    ADD = 1
    SUB = 2
    MUL = 3
    DIV = 4
    SQRT = 5
    EXP = 6
    LOG = 7


# 2. 定义操作到kernel模板的映射
op_templates = {
    # 二元操作
    Operation.ADD: lambda a, b, c: f"{c} = {a} + {b};",
    Operation.SUB: lambda a, b, c: f"{c} = {a} - {b};",
    Operation.MUL: lambda a, b, c: f"{c} = {a} * {b};",
    Operation.DIV: lambda a, b, c: f"{c} = {a} / {b};",
    # 一元操作
    Operation.SQRT: lambda a, c: f"{c} = sqrt({a});",
    Operation.EXP: lambda a, c: f"{c} = exp({a});",
    Operation.LOG: lambda a, c: f"{c} = log({a});",
}


# 3. 合并kernel生成器
def create_combined_kernel(operations, dependencies, output_var):
    """
    创建组合elementwise kernel

    参数:
        operations: 操作列表，每个元素是(操作类型, 输出变量名)
        dependencies: 依赖字典 {输出变量名: [输入变量名1, 输入变量名2]}
        output_var: 最终输出的变量名
    """
    # 收集所有变量
    all_vars = set()
    for _, out_var in operations:
        all_vars.add(out_var)
        all_vars.update(dependencies.get(out_var, []))

    # 确定输入变量（未被任何操作生成的变量）
    output_vars = {op[1] for op in operations}
    input_vars = sorted(all_vars - output_vars)  # 排序保证顺序一致

    # 生成内核参数声明
    input_decl = ", ".join([f"float32 {var}" for var in input_vars])
    output_decl = f"float32 {output_var}"

    # 生成内核代码
    code_lines = []

    # 声明中间变量
    for var in output_vars:
        if var != output_var:  # 最终输出由kernel返回
            code_lines.append(f"float {var};")

    # 添加操作代码
    for op_type, out_var in operations:
        inputs = dependencies[out_var]
        if len(inputs) == 1:
            code = op_templates[op_type](inputs[0], out_var)
        elif len(inputs) == 2:
            code = op_templates[op_type](inputs[0], inputs[1], out_var)
        else:
            raise ValueError("只支持一元或二元操作")
        code_lines.append(code)

    # 添加最终输出
    code_lines.append(f"{output_var} = {operations[-1][1]};")
    operation_code = "\n".join(code_lines)
    print(operation_code)
    # 创建组合kernel
    kernel = cp.ElementwiseKernel(
        input_decl, output_decl, operation_code, "combined_kernel"
    )

    return kernel, input_vars


# 4. 使用示例
if __name__ == "__main__":
    # 定义计算流程: (a + b) * exp(c) / sqrt(d)
    operations = [
        (Operation.ADD, "tmp1"),  # tmp1 = a + b
        (Operation.EXP, "tmp2"),  # tmp2 = exp(c)
        (Operation.MUL, "tmp3"),  # tmp3 = tmp1 * tmp2
        (Operation.SQRT, "tmp4"),  # tmp4 = sqrt(d)
        (Operation.DIV, "result"),  # result = tmp3 / tmp4
    ]

    # 定义依赖关系
    dependencies = {
        "tmp1": ["a", "b"],
        "tmp2": ["c"],
        "tmp3": ["tmp1", "tmp2"],
        "tmp4": ["d"],
        "result": ["tmp3", "tmp4"],
    }

    # 创建组合kernel
    kernel, input_order = create_combined_kernel(operations, dependencies, "result")

    # 生成测试数据
    a = cp.array([1.0, 2.0, 3.0], dtype=cp.float32)
    b = cp.array([4.0, 5.0, 6.0], dtype=cp.float32)
    c = cp.array([0.0, 1.0, 2.0], dtype=cp.float32)
    d = cp.array([1.0, 4.0, 9.0], dtype=cp.float32)

    # 按照input_order顺序准备输入
    inputs = [a, b, c, d]

    # 计算结果
    result = kernel(*inputs)

    # 验证结果
    expected = (a + b) * cp.exp(c) / cp.sqrt(d)
    print("计算结果:", result)
    print("预期结果:", expected)
    print("最大误差:", cp.abs(result - expected).max())
